package uk.ac.ebi.curtain.exec;

import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.atomic.AtomicReferenceArray;

import org.apache.commons.lang.StringUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import uk.ac.ebi.curtain.utils.Mapper;
import uk.ac.ebi.curtain.utils.data.DataEntry;
import uk.ac.ebi.curtain.utils.data.DataStore;
import uk.ac.ebi.curtain.utils.data.FileType;
import uk.ac.ebi.curtain.utils.data.IdentifierProcessor;
import uk.ac.ebi.curtain.utils.data.MetaData;
import uk.ac.ebi.curtain.utils.data.ReadType;
import uk.ac.ebi.curtain.utils.data.impl.DataStoreImpl;
import uk.ac.ebi.curtain.utils.data.impl.GenericDataHolder;
import uk.ac.ebi.curtain.utils.data.impl.SolexaIdentifierProcessor;
import uk.ac.ebi.curtain.utils.data.store.FileMetaStore;
import uk.ac.ebi.curtain.utils.data.store.SequenceMetaStore;
import uk.ac.ebi.curtain.utils.file.EntryMapperBean;
import uk.ac.ebi.curtain.utils.file.SequenceQualityType;
import uk.ac.ebi.curtain.utils.file.WriteUtil;
import uk.ac.ebi.curtain.utils.io.AbstractIO.CompressionFormat;
import uk.ac.ebi.curtain.utils.io.impl.FileIO;
import uk.ac.ebi.velvet.model.Strand;
import uk.co.flamingpenguin.jewel.cli.ArgumentValidationException;
import uk.co.flamingpenguin.jewel.cli.CliFactory;
import uk.co.flamingpenguin.jewel.cli.Option;

public class BulkFilterFastqByQualityScore {
	private static final Log log = LogFactory.getLog(BulkFilterFastqByQualityScore.class);
	private final AtomicInteger minLength = new AtomicInteger(30);
	private final AtomicInteger minQual = new AtomicInteger(30);
	private final AtomicInteger tileLen = new AtomicInteger(20);
	private final AtomicReference<ReadType> readType = new AtomicReference<ReadType>(ReadType.Short);
	private final AtomicReference<FileType> fileType = new AtomicReference<FileType>(FileType.fastq);
	private final AtomicReference<WriteUtil> writer = new AtomicReference<WriteUtil>(FileType.fastq.getWriter());
	private final AtomicReference<SequenceQualityType> inputType = new AtomicReference<SequenceQualityType>(SequenceQualityType.Sanger);
	private final AtomicReference<SequenceQualityType> targetType = new AtomicReference<SequenceQualityType>(SequenceQualityType.Sanger);
	private final AtomicReference<PrintWriter> errOut = new AtomicReference<PrintWriter>();
	private final AtomicReferenceArray<PrintWriter> outArr;
	private final AtomicReference<IdentifierProcessor> idProc = new AtomicReference<IdentifierProcessor>(new SolexaIdentifierProcessor());
	private final AtomicLong idCount = new AtomicLong(0);
	private final AtomicReference<String> prefix = new AtomicReference<String>(StringUtils.EMPTY);
	
	public BulkFilterFastqByQualityScore(int outSize) {
		this.outArr =  new AtomicReferenceArray<PrintWriter>(outSize);
	}
	
	public void setFileType(FileType fileType) {
		this.fileType.set(fileType);
		this.writer.set(fileType.getWriter());
	}

	private static interface FilterFastqByQualityScope {
		
		@Option(shortName = "d", description="Output directory. File name is gereated depending on the library.")
		public File getOutput();
		
		public boolean isOutput();
		// defaultValue=new String[]{SequenceQualityType.Sanger.toString()}
		@Option(shortName = "e", defaultValue="Illumina",description="Output Quality score encoding: Sanger, Illumina (1.3+)")
		public SequenceQualityType getMethod();
		
		public boolean isMethod();
		
		@Option(shortName="f",description="Property file")
		public File getFile();
		
		public boolean isFile();
		
		@Option(shortName="p",description="Paired file flag - if one file is proveded and this flag is set, the input file is understand to be shuffeled")
		public boolean getPaired();
		
		public boolean isPaired();
		
		@Option(shortName="l",description="Minimum read length, before excluded", defaultValue="30")
		public Integer getMinLength();
		
		@Option(shortName="q", description="Minimum average quality to keep over tile", defaultValue="30")
		public Integer getMinQuality();	
		
		@Option(shortName="t", description="Tile length to evaluate", defaultValue="20")
		public Integer getTileLength();
		
		@Option(shortName="z", description="Gzip output file(s)")
		public boolean getGzip();
		
		@Option(shortName="s",description="Write paired output into two (F/R) separate output files.")
		public boolean getSplitPaired();
		
		@Option(helpRequest = true, description = "display helpfull menu", shortName = "h")
		boolean getHelp();		
	}
	
	public static void main(String[] args) {
		FilterFastqByQualityScope param = null;
		try {
			param = CliFactory.parseArguments(FilterFastqByQualityScope.class, args);
		} catch (ArgumentValidationException e1) {
			System.err.println(e1.getMessage());
			System.exit(1);
		}
		if(!param.isOutput()){
			throw new IllegalArgumentException("Not output file(s) provided!!!");
		}
		if(!param.isFile()){
			throw new IllegalArgumentException("Property file required!!");
		}
		
		FileType ftype = FileType.fastq;
		File outDir = param.getOutput();
		
		SequenceMetaStore<File> store = new FileMetaStore();

		store.load(param.getFile());
			
		for(Integer id : store.getSetIds()){

			String baseId = "C"+id+"-";
			log.info("Process Store set "  + id + " ... ");
			boolean paired = false;
			SequenceQualityType sType = null;
			for(String annot : store.getAnnotation(id)){
				if(annot.equalsIgnoreCase("ShortPaired")){
					paired = true;
				}
				if(annot.equalsIgnoreCase(SequenceQualityType.Illumina.toString()) 
						|| annot.equalsIgnoreCase(SequenceQualityType.Sanger.toString())
						){
					sType = SequenceQualityType.valueOf(annot);
				}
			}
			int outLen = 1;
			if(paired && param.getSplitPaired()){
				outLen = 2;
			}
			BulkFilterFastqByQualityScore filter = new BulkFilterFastqByQualityScore(outLen);
			filter.prefix.set(baseId);
			filter.minQual.set(param.getMinQuality());
			filter.minLength.set(param.getMinLength());
			filter.tileLen.set(param.getTileLength());
			if(param.isMethod()){
				filter.targetType.set(param.getMethod());
			}
			filter.setFileType(ftype);
			if(null != sType) {
				filter.inputType.set(sType);
			}
			filter.readType.set(paired?ReadType.ShortPaired: ReadType.Short);
			try {
				log.info("Write output to: " );
				File errFile = new File(outDir,"cat"+id+".err.fastq");
				log.info("Error Fastq entry file: " + errFile);
				filter.errOut.set(new FileIO(errFile).getPrintWriter());
				for(int i = 0; i < outLen; ++i){
					StringBuilder sb = new StringBuilder();
					sb.append("cat").append(id);
					if(outLen == 2){
						sb.append("_");
						if(i == 0){
							sb.append(Strand.FWD.getSolexaExt());
						} else {
							sb.append(Strand.REV.getSolexaExt());
						}
					}
					sb.append(".");
					if(paired){
						sb.append("paired");
					}else{
						sb.append("single");						
					}
					sb.append(ftype.getPostfix());
					if(param.getGzip()){
						sb.append(".");
						sb.append(CompressionFormat.Gzip.getDefaultSuffix());
					}
					File file = new File(outDir,sb.toString());
					log.info("Outfile " + i + ": " + file );
					filter.setOutput(i, new FileIO(file,(param.getGzip()?false:true)).getPrintWriter());
				}
				for(File[] arr : store.getEntries(id)){
					filter.process(arr);
				}
			} catch (IOException e) {
				e.printStackTrace();
			}finally{
				FileIO.closeQuietly(filter.setErrOut(null));
				for(int i = 0; i < outLen; ++i){
					FileIO.closeQuietly(filter.setOutput(i, null));
				}
			}
		}	
		log.info("Finished processing!!!!");
	}
	
	public PrintWriter setOutput(int i, PrintWriter pw){
		return this.outArr.getAndSet(i, pw);
	}
	
	public void setMinQual(Integer minQual) {
		this.minQual.set(minQual);
	}
	
	public void setMinLength(Integer minLen) {
		this.minLength.set(minLen);
	}
	
	public void setTileLen(Integer tileLen) {
		this.tileLen.set(tileLen);
	}
	
	public void setInputType(SequenceQualityType inputType) {
		this.inputType.set(inputType);
	}
	
	public void setTargetType(SequenceQualityType targetType) {
		this.targetType.set(targetType);
	}
	
	public FileType getFileType() {
		return fileType.get();
	}
	
	public ReadType getReadType() {
		return readType.get();
	}
	
	public IdentifierProcessor getIdProcessor() {
		return idProc.get();
	}
	
	public SequenceQualityType getInputType() {
		return inputType.get();
	}
	
	public SequenceQualityType getTargetType() {
		return targetType.get();
	}
	
	public WriteUtil getWriter() {
		return writer.get();
	}
	
	public String getPrefix() {
		return prefix.get();
	}
	
	private PrintWriter getErrOut() {
		return errOut.get();
	}
	
	public PrintWriter setErrOut(PrintWriter errOut) {
		return this.errOut.getAndSet(errOut);
	}
	
	public void process(File[] arr) {
		log.info("Process " + Arrays.toString(arr));
		DataEntry<File> data = new GenericDataHolder<File>(new MetaData(0, getFileType(), this.getReadType()), arr);
		DataStore<EntryMapperBean> dataStore = new DataStoreImpl(data);
		dataStore.parse(new Mapper<DataEntry<EntryMapperBean>>() {
			
			@Override
			public void map(DataEntry<EntryMapperBean> bean) {
				IdentifierProcessor proc = getIdProcessor();
				if(getReadType().isPaired()){
					if(bean.getEntries().length !=2){
						throw new IllegalStateException("Paired entries should always have two entries!!!" + bean.getEntries().length);
					}
					String fId = bean.getEntries()[0].getId();
					String rId = bean.getEntries()[1].getId();
					Strand fStr = proc.getStrand(fId);
					Strand rStr = proc.getStrand(rId);
					String cfId = proc.removeStrandInformation(fId);
					String crId = proc.removeStrandInformation(rId);
					if(!cfId.equals(crId)){
						throw new IllegalArgumentException(
								"Identifiers not the same for paired entry!!! " + cfId + " <> "+crId);
					}
					if(fStr.equals(rStr)){
						throw new IllegalStateException("Opposite strands expected!!! " + fId+ " and "+rId);
					}
				}
				if(trim(bean)){
					String id = getPrefix()+idCount.getAndIncrement();
					for(int i = 0; i < bean.getEntries().length; ++i){
						EntryMapperBean entry = bean.getEntries()[i];
						String sId = id; 
						if(getReadType().isPaired()){
							Strand strand = proc.getStrand(entry.getId());
							sId = proc.addStrandInformation(sId, strand);
						}
						entry.setHeader(entry.getId());
						entry.setId(sId);
						getWriter().write(outArr.get(i%outArr.length()), bean.getEntries()[i]);
					}
				} else {
					PrintWriter errPw = getErrOut();
					if(null != errPw){
						for(int i = 0; i < bean.getEntries().length; ++i){
							getWriter().write(errPw, bean.getEntries()[i]);
						}
					}
				}
			}
		});			
		log.info("Finished  " + Arrays.toString(arr));
	
	}

	private boolean trim(DataEntry<EntryMapperBean> bean){
		boolean valid = true;
		
		int arrLen = bean.getEntries().length;
		for(int i = 0; i < arrLen; ++i){
			EntryMapperBean entry = bean.getEntries()[i];
			if(entry.getSequenceLength() < getMinLength() ){
				valid = false;
				break;
			}
			byte[] qual = getInputType().toQuality(entry.getQuality());
			int start = getCutoffStart(entry,qual);
			int end = getCutoffEnd(entry,start,qual);
			int len = (end-start);
			if(len < getMinLength() ){
				valid = false;
				break;
			}else {
				doTrim(bean.getEntries()[i],start,end, qual,getTargetType());
			}
		}
		return valid;
	}

	private int getMinLength() {
		return minLength.get();
	}
	
	private void doTrim(EntryMapperBean bean, int start, int end, byte[] qual, SequenceQualityType to) {
		bean.setQuality(to.fromQuality(qual).substring(start, end));
		bean.setSequence(bean.getSequence().substring(start, end));
	}

	private int getCutoffStart(EntryMapperBean bean, byte[] qual){
		int wSize = 20;
		int len = qual.length;
		int end = len-wSize;
		int pos = 0;
		while(pos < end){
			byte[] window = Arrays.copyOfRange(qual, pos, pos+wSize);
			int sum = 0;
			for(int i = 0; i < wSize; ++i){
				sum += window[i];
			}
			int avg = sum/wSize;
			if(avg >= getMinQuality()){
				break;
			}
			++pos;
		}
		return pos;
	}

	private int getMinQuality() {
		return minQual.get();
	}
	
	private int getCutoffEnd(EntryMapperBean bean, int fStart, byte[] qual){
		int wSize = getTileLen();
		int len = qual.length;
		int end = len-wSize;
		int pos = fStart;
		while(pos < end){
			byte[] window = Arrays.copyOfRange(qual, pos, pos+wSize);
			int sum = 0;
			for(int i = 0; i < wSize; ++i){
				sum += window[i];
			}
			int avg = sum/wSize;
			if(avg < getMinQuality()){
				break;
			}
			++pos;
		}
		return pos+wSize;
	}

	private int getTileLen() {
		return tileLen.get();
	}
	
}
